import os
os.environ['CUDA_VISIBLE_DEVICES']='1'
import torch
from torchvision import models,transforms
from torch import nn, optim
import logging
from model import SiameseNetwork
from config import Config 
from loader import CustomDataset
import numpy as np
# from eveluate import Evaluator
import pdb
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch.nn.functional as F

# nohup python main.py >0601.logs 2>&1& 

logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def main(config):
    # 创建保存模型的目录
    if not os.path.isdir(config.model_path):
        os.mkdir(config.model_path)
    # 训练
    if torch.cuda.is_available():
        logger.info("GPU is avalible")
    else:
        logger.info('only use CPU')
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SiameseNetwork().to(device)
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters(),lr=config.learning_rate)

    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 适应 VGG16 输入
        transforms.ToTensor(),
    ])

    dataset = CustomDataset(root_dir=config.data_path, transform=transform)
    loader = torch.utils.data.DataLoader(dataset, config.batch_size)
    avg_train_losses = []
    for epoch in range(config.epoch):
        epoch += 1
        model.train()
        logger.info(f"epoch {epoch} begin")
        train_loss = []
        # pdb.set_trace()
        for index, batch_data in tqdm(enumerate(loader)):
            img1, img2, labels=batch_data[0][0],batch_data[0][1],batch_data[1]
            img1, img2, labels = img1.to(device), img2.to(device), labels.to(device)
            optimizer.zero_grad()
            out1, out2 = model(img1, img2)
            # 计算两个特征向量的相似度
            similarity = F.cosine_similarity(out1, out2)
            loss = criterion(similarity, labels)
            loss.backward()
            optimizer.step()

            train_loss.append(loss.item())
            # if index % int(len(loader)/2)==0:
            logger.info(f'batch loss {loss}')

        logger.info("epoch average loss: %f" % np.mean(train_loss))
        avg_train_loss = np.mean(train_loss)
        avg_train_losses.append(avg_train_loss)
        # acc = evaluator.eval(epoch)  # 训练一次进行一次评估
    plt.plot(range(1, config.epoch + 1), avg_train_losses)
    plt.xlabel('Epoch')
    plt.ylabel('Average Train Loss')
    plt.savefig(f'train_loss_{epoch}_cosine.png')

    model_path = os.path.join(config.model_path, "%s-epoch_%d.pth" % (config.model_type, epoch))
    torch.save(model.state_dict(), model_path)  
    return 


if __name__ == "__main__":
    # config=Config()
    # logger.info(config)
    config = Config()
    config.print_attributes()
    logger.info(f"last acc {main(config)},current config {config}")




